#!/bin/bash

# EnPlatform Assistant 超轻量级模型设置脚本
# 专门针对Excel处理和低资源环境优化

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() {
    echo -e "${BLUE}[MINIMAL-SETUP]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[MINIMAL-SETUP]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[MINIMAL-SETUP]${NC} $1"
}

log_error() {
    echo -e "${RED}[MINIMAL-SETUP]${NC} $1"
}

# 等待服务就绪
wait_for_service() {
    local service_name=$1
    local service_url=$2
    local max_attempts=20
    local attempt=1

    log_info "等待 $service_name 服务就绪..."
    
    while [ $attempt -le $max_attempts ]; do
        if curl -s "$service_url/api/tags" >/dev/null 2>&1; then
            log_success "$service_name 服务已就绪"
            return 0
        fi
        
        log_info "等待 $service_name 服务启动... ($attempt/$max_attempts)"
        sleep 15
        attempt=$((attempt + 1))
    done
    
    log_error "$service_name 服务启动超时"
    return 1
}

# 下载优化后的模型
download_minimal_models() {
    log_info "开始下载Excel处理优化模型..."
    log_info "基于验证结果，优先配置嵌入模型，对话模型作为备用"

    # 等待服务就绪
    wait_for_service "对话服务" "http://localhost:11434" || exit 1
    wait_for_service "嵌入服务" "http://localhost:8081" || exit 1
    
    log_info "=== 优先下载嵌入模型 (Excel处理核心) ==="

    # 下载已验证的高性能嵌入模型
    log_info "下载 all-MiniLM-L6-v2 (约90MB) - 已验证Excel处理专用嵌入模型"
    if docker exec enplatform-ollama-embedding ollama pull all-minilm:l6-v2; then
        log_success "✅ all-MiniLM-L6-v2 下载成功 - Excel嵌入功能已就绪"
    else
        log_warning "主要嵌入模型下载失败，尝试备用模型..."

        # 备用选择：nomic-embed-text
        log_info "下载 nomic-embed-text (约274MB) - 备用嵌入模型"
        if docker exec enplatform-ollama-embedding ollama pull nomic-embed-text; then
            log_success "nomic-embed-text 下载成功"
        else
            log_error "所有嵌入模型下载失败"
            return 1
        fi
    fi

    log_info "=== 下载对话模型 (备用功能) ==="
    log_warning "注意：基于M1 Mac验证结果，对话模型性能有限，建议主要使用嵌入+模板回答"

    # 按性能优先级下载对话模型
    log_info "下载 Qwen2:0.5B (约400MB) - 当前最佳ARM64轻量级模型"
    if docker exec enplatform-ollama-chat ollama pull qwen2:0.5b; then
        log_success "Qwen2:0.5B 下载成功"
    else
        log_warning "Qwen2 下载失败，尝试备用模型..."

        # 备用选择1：Gemma 2B
        log_info "下载 Gemma 2B (约1.4GB) - Google轻量级模型"
        if docker exec enplatform-ollama-chat ollama pull gemma:2b; then
            log_success "Gemma 2B 下载成功"
        else
            log_warning "Gemma 下载失败，尝试最后备用..."

            # 备用选择2：TinyLlama
            log_info "下载 TinyLlama 1.1B (约700MB) - 最后备用模型"
            if docker exec enplatform-ollama-chat ollama pull tinyllama:1.1b; then
                log_success "TinyLlama 1.1B 下载成功"
            else
                log_error "所有对话模型下载失败，但嵌入功能仍可正常使用"
            fi
        fi
    fi
    

    
    # 验证模型下载
    log_info "=== 验证模型下载 ==="
    
    log_info "=== 验证模型安装 ==="

    log_info "嵌入服务模型列表 (Excel处理核心):"
    docker exec enplatform-ollama-embedding ollama list

    log_info "对话服务模型列表 (备用功能):"
    docker exec enplatform-ollama-chat ollama list

    log_success "Excel处理模型配置完成！"
    log_info "推荐使用方式："
    log_info "  ✅ 主要功能: 嵌入模型进行Excel向量化和匹配"
    log_info "  ⚠️  备用功能: 对话模型性能有限，建议使用模板回答"
    log_info "总内存使用预计: 嵌入服务 ~200MB + 对话服务 ~800MB = ~1GB"
}

# 测试Excel处理功能
test_excel_processing() {
    log_info "=== 测试Excel处理功能 ==="
    
    # 检测可用的嵌入模型（优先级最高，Excel处理核心）
    EMBED_MODEL=""
    if docker exec enplatform-ollama-embedding ollama list | grep -q "all-minilm:l6-v2"; then
        EMBED_MODEL="all-minilm:l6-v2"
    elif docker exec enplatform-ollama-embedding ollama list | grep -q "nomic-embed-text"; then
        EMBED_MODEL="nomic-embed-text"
    else
        log_error "没有找到可用的嵌入模型 - Excel处理功能不可用"
        return 1
    fi

    # 检测可用的对话模型（备用功能）
    CHAT_MODEL=""
    if docker exec enplatform-ollama-chat ollama list | grep -q "qwen2:0.5b"; then
        CHAT_MODEL="qwen2:0.5b"
    elif docker exec enplatform-ollama-chat ollama list | grep -q "gemma:2b"; then
        CHAT_MODEL="gemma:2b"
    elif docker exec enplatform-ollama-chat ollama list | grep -q "tinyllama:1.1b"; then
        CHAT_MODEL="tinyllama:1.1b"
    else
        log_warning "没有找到可用的对话模型，但嵌入功能仍可正常使用"
        CHAT_MODEL=""
    fi
    
    log_info "使用嵌入模型: $EMBED_MODEL (Excel处理核心)"
    if [ -n "$CHAT_MODEL" ]; then
        log_info "使用对话模型: $CHAT_MODEL (备用功能)"
    else
        log_warning "对话模型不可用，将跳过对话测试"
    fi

    # 测试嵌入功能（最重要）
    log_info "1. 测试Excel数据嵌入功能 (核心功能)..."
    
    local test_texts=(
        "TEMP_01 环境温度 float 量测属性"
        "STATUS_01 设备状态 enum 量测属性"
        "CTRL_01 控制指令 bool 控制属性"
    )
    
    for text in "${test_texts[@]}"; do
        log_info "测试文本: $text"
        
        local result=$(curl -s -X POST http://localhost:8081/api/embeddings \
            -H "Content-Type: application/json" \
            -d "{\"model\": \"$EMBED_MODEL\", \"prompt\": \"$text\"}")
        
        if echo "$result" | grep -q "embedding"; then
            local dim=$(echo "$result" | python3 -c "import sys, json; data=json.load(sys.stdin); print(len(data['embedding']) if 'embedding' in data else 0)" 2>/dev/null || echo "0")
            if [ "$dim" -gt 0 ]; then
                log_success "嵌入成功，维度: $dim"
            else
                log_warning "嵌入响应格式异常"
            fi
        else
            log_warning "嵌入失败: $text"
        fi
    done
    
    # 测试对话功能（如果可用）
    if [ -n "$CHAT_MODEL" ]; then
        log_info "2. 测试Excel数据对话功能 (备用功能)..."
        log_warning "注意：基于验证结果，对话功能性能有限，建议使用嵌入+模板回答"

        local excel_question="分析这个Excel数据：TEMP_01 环境温度 float 量测属性。这是什么类型的数据？"

        log_info "问题: $excel_question"

        # 设置较短的超时时间
        local response=$(timeout 15s curl -s -X POST http://localhost:11434/api/generate \
            -H "Content-Type: application/json" \
            -d "{
                \"model\": \"$CHAT_MODEL\",
                \"prompt\": \"$excel_question\",
                \"stream\": false,
                \"options\": {
                    \"temperature\": 0.7,
                    \"num_predict\": 100
                }
            }" 2>/dev/null || echo '{"error": "timeout"}')

        if echo "$response" | grep -q "response"; then
            local answer=$(echo "$response" | python3 -c "import sys, json; data=json.load(sys.stdin); print(data.get('response', 'No response'))" 2>/dev/null || echo "解析失败")
            if [ "$answer" != "No response" ] && [ "$answer" != "解析失败" ]; then
                log_success "对话测试成功"
                echo "回答: $answer"
            else
                log_warning "对话响应解析失败"
            fi
        else
            log_warning "对话测试超时或失败 - 这是已知问题，建议使用嵌入+模板方案"
        fi
    else
        log_info "2. 跳过对话测试 (模型不可用)"
        log_info "推荐使用嵌入模型 + 模板回答的方案"
    fi
    
    # 测试Excel内容匹配
    log_info "3. 测试Excel内容匹配..."
    
    # 创建测试向量数据
    local search_query="温度传感器"
    log_info "搜索查询: $search_query"
    
    local query_embedding=$(curl -s -X POST http://localhost:8081/api/embeddings \
        -H "Content-Type: application/json" \
        -d "{\"model\": \"$EMBED_MODEL\", \"prompt\": \"$search_query\"}")
    
    if echo "$query_embedding" | grep -q "embedding"; then
        log_success "查询嵌入生成成功"
        log_info "Excel内容匹配功能基础验证通过"
    else
        log_warning "查询嵌入生成失败"
    fi
    
    log_success "=== Excel处理功能测试完成 ==="
}

# 显示资源使用情况
show_resource_usage() {
    log_info "=== 系统资源使用情况 ==="
    
    echo "Docker容器资源使用:"
    docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.MemPerc}}" | grep enplatform
    
    echo ""
    echo "模型文件大小:"
    docker exec enplatform-ollama-minimal du -sh /root/.ollama/models/* 2>/dev/null | head -5 || echo "无法获取对话模型大小"
    docker exec enplatform-embedding-minimal du -sh /root/.ollama/models/* 2>/dev/null | head -5 || echo "无法获取嵌入模型大小"
}

# 主函数
main() {
    case "${1:-setup}" in
        setup)
            download_minimal_models
            ;;
        test)
            test_excel_processing
            ;;
        resources)
            show_resource_usage
            ;;
        *)
            echo "EnPlatform Assistant 超轻量级模型设置"
            echo ""
            echo "用法: $0 [命令]"
            echo ""
            echo "命令:"
            echo "  setup      下载超轻量级模型"
            echo "  test       测试Excel处理功能"
            echo "  resources  显示资源使用情况"
            ;;
    esac
}

main "$@"
